package middleware

import (
	"crypto/subtle"
	"math/rand"
	"net/http"
	
	"github.com/gin-gonic/gin"
)

const csrfTokenLength = 32
const csrfTokenCookie = "csrf_token"
const csrfTokenHeader = "X-CSRF-Token"

// CSRFMiddleware 生成和验证CSRF令牌
func CSRFMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 对于GET请求，生成CSRF令牌
		if c.Request.Method == http.MethodGet {
			token := generateCSRFToken(csrfTokenLength)
			c.SetCookie(csrfTokenCookie, token, 3600, "/", "", false, true)
			c.Set("csrf_token", token)
			c.Next()
			return
		}
		
		// 对于修改请求，验证CSRF令牌
		cookieToken, _ := c.Cookie(csrfTokenCookie)
		headerToken := c.GetHeader(csrfTokenHeader)
		
		// 使用恒定时间比较防止计时攻击
		if subtle.ConstantTimeCompare([]byte(cookieToken), []byte(headerToken)) != 1 {
			c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "CSRF token mismatch"})
			return
		}
		
		c.Next()
	}
}

// 生成随机CSRF令牌
func generateCSRFToken(length int) string {
	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	b := make([]byte, length)
	for i := range b {
		b[i] = charset[rand.Intn(len(charset))]
	}
	return string(b)
}